from datetime import datetime
import backtrader as bt
import akshare as ak
import pandas as pd
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks(index):
    sql = f"SELECT DISTINCT CONSTITUENT_CODE FROM index_constituent WHERE INDEX_CODE='{index}'"
    r = exec_sql(sql)
    stocks = [tok[0] for tok in r]
    # stocks = [tok[0] for tok in r if tok[0][:2] not in ('30', '68')]
    # stocks = ['300122', '600809']

    return stocks


def get_data(stocks, freq='day'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'HIS_DAILY_DATA'  # HIS_DAILY_DATA

    test_code = '399300'
    # stocks.append('000877')
    stocks.insert(0, test_code)
    # stocks = ['399300', '000001', '001289', '000063']
    stocks_str = ','.join(f"'{tok}'" for tok in stocks)
    sql = f"SELECT code, datetime, [close], [open], low, high, vol, amount FROM {table_name} WHERE code in ({stocks_str})"
    # sql = f"SELECT * FROM {table_name} WHERE code in ({stocks_str})"
    r = exec_sql(sql)
    # df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume', 'money'])
    df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume', 'money'])
    df['datetime'] = df['datetime'].map(time_map)
    # sql2 = f"SELECT COUNT(1) FROM {table_name} WHERE code='{test_code}'"
    # r2 = exec_sql(sql2)
    # test_count = r2[0][0]
    res = dict()

    for code in stocks:
        sig_df = df.loc[df['code'] == code]
        sig_df.sort_values('datetime', inplace=True, ignore_index=True)
        sig_df.set_index('datetime', inplace=True)
        res[code] = sig_df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y%m%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


pre_ratio = 0.0048
ratio = 0.1
trail_ratio = 0.035
count_info = dict()
# stocks = get_stocks()
init_cash = 1e6


def run_backtrade(stocks, index):
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TestStrategy)
    data_all = get_data(stocks)
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    for i, code in enumerate(data_all):
        if code == index:
            continue
        print(f"i={i}->{code} Done !")
        data = data_all[code]

        datafeed = bt.feeds.PandasData(dataname=data)
        # st_date = datetime(2023, 1, 1)
        # ed_date = datetime(2023, 6, 16)
        # datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=ed_date)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(f'returns={returns}')
    qs.reports.html(returns, title='macd_kdj', output='stats.html', download_filename='stats.html')
    qs.reports.basic(returns)
    f = webbrowser.open('stats.html')
    # pf.create_full_tear_sheet(returns)
    # cerebro.plot()


class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        ('money', 1e4),
        ('printlog', False),
    )

    def __init__(self):
        self.order = None
        self.daily_pnl = dict()
        self.trade_info = list()
        self.signal = dict()
        self.sell_signal = dict()
        self.buy_info = dict()
        self.buyprice = None
        self.buycomm = None

        for data in self.datas:
            # 9个交易日内最高价
            self.high_nine = bt.indicators.Highest(self.data.high, period=9, plot=False)
            # 9个交易日内最低价
            self.low_nine = bt.indicators.Lowest(self.data.low, period=9, plot=False)
            # 计算rsv值
            self.rsv = 100 * bt.DivByZero(
                self.data - self.low_nine, self.high_nine - self.low_nine, zero=None
            )
            # 计算rsv的3周期加权平均值，即K值
            self.K = bt.indicators.EMA(self.rsv, period=3, plot=False)
            # D值=K值的3周期加权平均值
            self.D = bt.indicators.EMA(self.K, period=3, plot=False)
            # J=3*K-2*D
            self.J = 3 * self.K - 2 * self.D
            self.sell_signal[data] = bt.ind.CrossOver(self.J, self.D)
            # bt.indicator.LinePlotterIndicator(self.sell_signal[data], name='sell_signal')

            # MACD策略参数
            dif = bt.indicators.EMA(data, period=12, plot=False) - bt.indicators.EMA(data, period=26, plot=False)
            dea = bt.indicators.EMA(dif, period=9)
            # self.signal[data] = dif - dea
            self.signal[data] = bt.ind.CrossOver(dif, dea)
            # bt.indicator.LinePlotterIndicator(self.signal[data], name='buy_signal')

    def downcast(self, price):
        res = self.p.money / price // 100 * 100
        return res

    def prenext(self):
        self.next()

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.value:.2f}, Size: {order.executed.size:.2f}')

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.buy_info[code] = order.executed.price
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.value:.2f}, Size: {order.executed.size:.2f}')
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def next(self):
        if self.order:
            return

        for data in self.datas:
            code = data._name
            if code == '399300':
                continue
            dt = data.datetime.date(0)  # 获取当前的回测时间点
            if dt != self.datas[0].datetime.date(0):
                continue
            dt = int(dt.strftime('%Y%m%d'))

            if dt not in self.daily_pnl:
                # item = [dt, self.broker.getvalue() - init_cash]
                self.daily_pnl[dt] = round(self.broker.getvalue() - init_cash, 2)

            pos = self.getposition(data).size
            self.log(f'code={code}->{data.datetime.datetime(0)}->pos={pos}')

            if not pos:
                if self.signal[data] > 0:
                    self.log(f'code={code}->{data.datetime.datetime(0)}->signal={self.signal[data][0]}')
                    amount = self.downcast(data.close[0])
                    self.order = self.buy(data=data, size=amount, name=code)

            elif code in self.buy_info:
                condition = data.close[0] / self.buy_info[code] - 1
                self.log(
                    f'code={code}->{data.datetime.datetime(0)}->close={data.close[0]}->cost={self.buy_info[code]}->signal={self.sell_signal[data][0]}->condition={condition}')
                if (self.sell_signal[data] < 0) and (condition > 0.1 or condition < -0.1):
                    self.order = self.close(data=data, name=code)
                    # del self.buy_info[code]

    def stop(self):
        print('Strategy Finish!!!')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)

        daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
        append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                 columns=['DATE_T', 'PNL'])
        daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
        trade_df = pd.DataFrame(self.trade_info, columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
        pnl_table = 'TEST_DAILY_PNL'
        trade_table = 'TEST_TRADE_INFO'
        # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
        insert_db(daily_pnl_df, pnl_table)
        insert_db(trade_df, trade_table)
        print('insert table ok')


def main():
    index = '399300'
    stocks = get_stocks(index)
    run_backtrade(stocks, index)


if __name__ == '__main__':
    main()
